import pandas as pd
from pathlib import Path
import os
import traceback

# 配置数据文件路径
DATA_DIR = Path('data')
SOURCE_DIR = DATA_DIR / 'source'
OUTPUT_DIR = DATA_DIR / 'output'

def get_relative_relationships(folder_path):
    """
    从关系总表中获取亲属关系数据
    只获取关系类型为'亲属'的行的名称和关系列
    """
    try:
        # 从output目录读取关系总表
        output_folder = OUTPUT_DIR / folder_path.name
        relationship_file = output_folder / "关系总表.xlsx"
        if not relationship_file.exists():
            print(f"[DEBUG] 关系总表文件不存在: {relationship_file}")
            return None
            
        # 读取Excel文件
        df = pd.read_excel(relationship_file)
        
        # 检查必要的列是否存在
        required_columns = ['名称', '关系', '关系类型']
        if not all(col in df.columns for col in required_columns):
            print(f"[DEBUG] 关系总表缺少必要的列: {[col for col in required_columns if col not in df.columns]}")
            return None
            
        # 筛选关系类型为"亲属"的数据
        relatives_df = df[df['关系类型'] == '亲属'].copy()
        if relatives_df.empty:
            print(f"[DEBUG] 关系总表中未找到亲属关系数据")
            return None
            
        # 去除名称为空的行
        relatives_df = relatives_df[relatives_df['名称'].notna()]
        # 去除重复值
        relatives_df = relatives_df.drop_duplicates(subset=['名称'])
        # 添加关系类型和公司名称列（保持与其他关系数据格式一致）
        relatives_df['关系类型'] = '亲属'
        relatives_df['公司名称'] = None
        
        return relatives_df[['名称', '关系', '关系类型', '公司名称']]
    except Exception as e:
        print(f"[DEBUG] 从关系总表获取亲属关系数据时出错: {str(e)}")
        traceback.print_exc()
    return None

def get_employment_relationships(folder_path):
    """获取任职单位关系数据"""
    try:
        employment_file = folder_path / "任职单位.xlsx"
        if not employment_file.exists():
            print(f"[DEBUG] 任职单位文件不存在: {employment_file}")
            return None
            
        df = pd.read_excel(employment_file)
        if '任职单位' not in df.columns:
            print(f"[DEBUG] 任职单位文件缺少必要的列: 任职单位")
            return None
            
        # 创建关系数据
        result_df = pd.DataFrame()
        result_df['名称'] = df['任职单位']
        result_df['关系'] = '任职单位'
        result_df['关系类型'] = '任职单位'
        result_df['公司名称'] = None
        return result_df
    except Exception as e:
        print(f"[DEBUG] 处理任职单位文件出错: {str(e)}")
        traceback.print_exc()
    return None

def get_specific_relationships(folder_path):
    """获取特定关系人数据"""
    try:
        specific_file = folder_path / "特定关系人.xlsx"
        if not specific_file.exists():
            print(f"[DEBUG] 特定关系人文件不存在: {specific_file}")
            return None
            
        df = pd.read_excel(specific_file)
        if '名称' not in df.columns:
            print(f"[DEBUG] 特定关系人文件缺少必要的列: 名称")
            return None
            
        # 创建关系数据
        result_df = df[['名称']].copy()
        if '关系' in df.columns:
            result_df['关系'] = df['关系']
        else:
            result_df['关系'] = '特定关系人'
        result_df['关系类型'] = '特定关系人'
        result_df['公司名称'] = None
        return result_df
    except Exception as e:
        print(f"[DEBUG] 处理特定关系人文件出错: {str(e)}")
        traceback.print_exc()
    return None

def get_san_shang_relationships(folder_path):
    """获取三商清单关系数据"""
    try:
        output_folder = OUTPUT_DIR / folder_path.name
        san_shang_file = output_folder / "三商清单处理结果.xlsx"
        if not san_shang_file.exists():
            print(f"[DEBUG] 三商清单处理结果文件不存在: {san_shang_file}")
            return None
            
        df = pd.read_excel(san_shang_file)
        if not all(col in df.columns for col in ['联系人', '身份', '公司名称']):
            print(f"[DEBUG] 三商清单处理结果文件缺少必要的列")
            return None
            
        # 创建关系数据
        result_df = pd.DataFrame()
        result_df['名称'] = df['联系人']
        result_df['关系'] = df['身份']
        result_df['关系类型'] = '三商清单'
        result_df['公司名称'] = df['公司名称']
        return result_df
    except Exception as e:
        print(f"[DEBUG] 处理三商清单文件出错: {str(e)}")
        traceback.print_exc()
    return None

def generate_relationship_table():
    """
    生成关系总表
    为每个文件夹生成关系总表，包含该文件夹中的所有关系数据（亲属、任职单位、特定关系人、三商清单）
    """
    try:
        print("\n[DEBUG] 开始生成关系总表")
        
        # 确保目录存在
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        SOURCE_DIR.mkdir(parents=True, exist_ok=True)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
        
        # 遍历source目录下的所有文件夹
        for folder in SOURCE_DIR.iterdir():
            if folder.is_dir():
                folder_name = folder.name
                print(f"[DEBUG] 处理文件夹: {folder_name}")
                
                # 收集所有关系数据
                relationship_dfs = []
                
                # 获取各类关系数据
                relative_df = get_relative_relationships(folder)
                if relative_df is not None:
                    relationship_dfs.append(relative_df)
                    print(f"[DEBUG] 获取到 {len(relative_df)} 条亲属关系数据")
                
                employment_df = get_employment_relationships(folder)
                if employment_df is not None:
                    relationship_dfs.append(employment_df)
                    print(f"[DEBUG] 获取到 {len(employment_df)} 条任职单位关系数据")
                
                specific_df = get_specific_relationships(folder)
                if specific_df is not None:
                    relationship_dfs.append(specific_df)
                    print(f"[DEBUG] 获取到 {len(specific_df)} 条特定关系人数据")
                
                san_shang_df = get_san_shang_relationships(folder)
                if san_shang_df is not None:
                    relationship_dfs.append(san_shang_df)
                    print(f"[DEBUG] 获取到 {len(san_shang_df)} 条三商清单关系数据")
                
                if relationship_dfs:
                    # 合并所有关系数据
                    combined_df = pd.concat(relationship_dfs, ignore_index=True)
                    # 去除重复数据
                    combined_df = combined_df.drop_duplicates()
                    # 添加被审查人列
                    combined_df['被审查人'] = folder_name
                    
                    # 确保对应的output文件夹存在
                    output_folder = OUTPUT_DIR / folder_name
                    output_folder.mkdir(parents=True, exist_ok=True)
                    
                    # 保存到对应文件夹中的Excel文件
                    output_file = output_folder / "关系总表.xlsx"
                    combined_df.to_excel(output_file, index=False)
                    print(f"[DEBUG] 关系总表已保存: {output_file}")
                    print(f"[DEBUG] 共写入 {len(combined_df)} 条数据")
                else:
                    print(f"[DEBUG] {folder_name} 未找到任何关系数据")
        
        return True
        
    except Exception as e:
        print(f"[DEBUG] 生成关系总表时出错: {str(e)}")
        traceback.print_exc()
        return False 